"""
内容：输入视频预测多元分类问题的模型示例
日期：2020年7月6日
作者：Howie
"""

# 调用
import numpy as np
import matplotlib.pyplot as plt
from keras.utils import plot_model, to_categorical
from keras.models import Sequential
from keras.layers import Dense, Conv2D, MaxPooling2D, Flatten, Dropout
from keras.callbacks import EarlyStopping

# 预设
DATASET_PATH = '../dataset/mnist/mnist.npz'  # 数据集路径
PLT_ROW = 5
PLT_COL = 5
# 超参
N_SAMPLES_TRAIN = 60000     # 训练集样本数量
N_SAMPLES_TEST = 10000  # 测试集样本数量
IMG_SIZE = 28  # 图片尺寸
N_CLS = 10  # 类别数
EPOCHS = 100  # 周期
BATCH_SIZE = 32  # 批次数量
VAL_SET_SPLIT = 0.2  # 验证集比例
early_stopping = EarlyStopping(  # 设置早停，监控每个训练周期的验证精度
    monitor='val_accuracy',
    patience=8
)
CALL_BACK_FUNC = [early_stopping]  # 回调函数


def hist_plot(hist, model_name):
    """
    # 可视化训练过程
    :param hist: history对象
    :return:
    """
    fig, loss_ax = plt.subplots()
    acc_ax = loss_ax.twinx()
    # 每个训练周期的训练误差与验证误差
    loss_ax.plot(hist.history['loss'], 'y', label='Loss')
    loss_ax.plot(hist.history['val_loss'], 'r', label='Val Loss')
    loss_ax.set_ylim([0.0, 0.5])
    # 每个训练周期的训练精度与验证精度
    acc_ax.plot(hist.history['accuracy'], 'b', label='Acc')
    acc_ax.plot(hist.history['val_accuracy'], 'g', label='Val Acc')
    acc_ax.set_ylim([0.8, 1.0])
    # 横轴与纵轴
    loss_ax.set_xlabel('epoch')
    loss_ax.set_ylabel('loss')
    acc_ax.set_ylabel('accuracy')
    # 标签
    loss_ax.legend(loc='upper left')
    acc_ax.legend(loc='lower left')
    # 标题
    plt.title(model_name)
    # 保存
    plt.savefig('./logs/History_Demo6_' + model_name + '.pdf')
    # 展示
    plt.show()


def load_dataset():
    """
    # 生成数据集
    :return: 划分好的训练集和测试集
    """
    with np.load(DATASET_PATH) as data:
        (X_train, Y_train), (X_test, Y_test) = (
            data['x_train'], data['y_train']), (data['x_test'], data['y_test'])
    # 独热编码处理
    Y_train = to_categorical(Y_train)
    Y_test = to_categorical(Y_test)
    # 用于MLP
    X_train_1D = X_train.reshape(
        N_SAMPLES_TRAIN,
        IMG_SIZE * IMG_SIZE).astype('float32') / 255.0  # 数据正则化
    X_test_1D = X_test.reshape(N_SAMPLES_TEST,
                               IMG_SIZE * IMG_SIZE).astype('float32') / 255.0
    # 用于CNN
    X_train = X_train.reshape(
        N_SAMPLES_TRAIN,
        IMG_SIZE,
        IMG_SIZE,
        1).astype('float32') / 255.0
    X_test = X_test.reshape(
        N_SAMPLES_TEST,
        IMG_SIZE,
        IMG_SIZE,
        1).astype('float32') / 255.0
    print('train samples: {}\n'
          'test samples: {}'.format(X_train.shape[0], X_test.shape[0]))

    # 多元分类问题中的标签值指定为0~9，但此处改为奇数/偶数的二元分类标签值
    # 1代表奇数，0代表偶数
    return X_train, X_train_1D, Y_train, X_test, X_test_1D, Y_test


def dataset_visualization(sample, label):
    """
    将生成的数据集的一部分进行可视化展示
    :return:
    """
    plt.rcParams["figure.figsize"] = (10, 10)
    fig, ax = plt.subplots(PLT_ROW, PLT_COL)

    for i in range(PLT_ROW * PLT_COL):
        sub_plt = ax[i // PLT_ROW, i % PLT_ROW]
        sub_plt.axis('off')
        sub_plt.imshow(sample[i].reshape(IMG_SIZE, IMG_SIZE))
        sub_plt_title = 'R: ' + str(np.argmax(label[i]))
        sub_plt.set_title(sub_plt_title)
    plt.savefig('./logs/Dataset Visualization_Demo6.pdf')
    plt.show()


def model_building(choice='MLP'):
    """
    搭建模型
    :return:
    """
    model = Sequential()
    if choice == 'MLP':  # 多层感知器神经网络模型
        model.add(Dense(256, input_dim=IMG_SIZE * IMG_SIZE, activation='relu'))
        model.add(Dense(256, input_dim=IMG_SIZE * IMG_SIZE, activation='relu'))
        model.add(Dense(256, input_dim=IMG_SIZE * IMG_SIZE, activation='relu'))
        model.add(Dense(N_CLS, activation='softmax'))

    elif choice == 'CNN':  # 卷积神经网络模型
        model.add(
            Conv2D(
                32, (3, 3), activation='relu', input_shape=(
                    IMG_SIZE, IMG_SIZE, 1)))
        model.add(MaxPooling2D(pool_size=(2, 2)))
        model.add(Conv2D(32, (3, 3), activation='relu'))
        model.add(MaxPooling2D(pool_size=(2, 2)))
        model.add(Flatten())
        model.add(Dense(256, activation='relu'))
        model.add(Dense(N_CLS, activation='softmax'))
    elif choice == 'DCNN':  # 深度卷积神经网络模型
        model.add(
            Conv2D(
                32, (3, 3), activation='relu', input_shape=(
                    IMG_SIZE, IMG_SIZE, 1)))
        model.add(Conv2D(32, (3, 3), activation='relu'))
        model.add(MaxPooling2D(pool_size=(2, 2)))
        model.add(Dropout(0.25))
        model.add(Conv2D(64, (3, 3), activation='relu'))
        model.add(Conv2D(64, (3, 3), activation='relu'))
        model.add(MaxPooling2D(pool_size=(2, 2)))
        model.add(Dropout(0.25))
        model.add(Flatten())
        model.add(Dense(256, activation='relu'))
        model.add(Dropout(0.5))
        model.add(Dense(N_CLS, activation='softmax'))
    else:
        print("Please choose a model!")
    # 可视化模型结构
    plot_model(model, to_file='./logs/Demo6_model.pdf', show_shapes=True)
    # 返回搭建好的模型
    return model


def model_training():
    """
    # 训练模型
    :return:
    """
    # 准备数据集
    X_train, X_train_1D, Y_train, X_test, X_test_1D, Y_test = load_dataset()
    dataset_visualization(X_train_1D, Y_train)
    choices = ['MLP', 'CNN', 'DCNN']
    for choice in choices:
        # 选择模型
        model = model_building(choice=choice)
        # 设置模型训练过程
        model.compile(
            loss='categorical_crossentropy',
            optimizer='sgd',
            metrics=['accuracy'])
        # 训练过程
        hist = model.fit(
            X_train_1D if choice == 'MLP' else X_train,
            Y_train,
            epochs=EPOCHS,
            batch_size=BATCH_SIZE,
            validation_split=VAL_SET_SPLIT,
            callbacks=CALL_BACK_FUNC)
        # 查看训练过程
        print(hist.history)
        hist_plot(hist, model_name=choice)
        # 评价模型
        model_evaluation(model, X_test_1D if choice ==
                         'MLP' else X_test, Y_test)
        model_application(model, X_test_1D if choice ==
                          'MLP' else X_test, Y_test)


def model_evaluation(model, X_test, Y_test):
    """
    # 评价模型
    :param model: 模型
    :param X_test: 测试集样本
    :param Y_test: 测试集标签
    :return:
    """
    loss_and_metrics = model.evaluate(X_test, Y_test, batch_size=BATCH_SIZE)
    print("Evaluation loss and metrics: {}".format(loss_and_metrics))


def model_application(model, X_test, Y_test):
    """
    # 调用模型
    :param model: 模型
    :param X_test: 测试集样本
    :param Y_test: 测试集标签
    :return:
    """
    Y_hat_test = model.predict(X_test, batch_size=BATCH_SIZE)
    # 预测结果可视化
    plt.rcParams["figure.figsize"] = (10, 10)
    fig, ax = plt.subplots(PLT_ROW, PLT_COL)

    cnt, i = 0, 0
    for i in range(PLT_ROW * PLT_COL):

        sub_plt = ax[cnt // PLT_ROW, cnt % PLT_COL]
        sub_plt.axis('off')
        sub_plt.imshow(X_test[i].reshape(IMG_SIZE, IMG_SIZE))

        sub_plt_title = 'R: ' + str(np.argmax(Y_test[i])) + 'P: ' + str(np.argmax(Y_hat_test[i]))

        sub_plt.set_title(sub_plt_title)

        i += 1
        cnt += 1

    plt.savefig('./logs/Prediction_Demo6.pdf')
    plt.show()


if __name__ == '__main__':
    model_training()
